#-------------------------------------------------------------------------------
# Name:        kerf-ucb
# Purpose:     Split phylo trees based on percent similarity
#
# Author:      lgcarter
#
# Created:     04/10/2011
# Copyright:   (c) lgcarter 2011
#-------------------------------------------------------------------------------
#!/usr/bin/env python

from Bio import Phylo
from cStringIO import StringIO
import sys

def main(cutoff, align_file, tree_file):
    """ reads cutoff value, and alignment and tree text files and returns csv and individual a2m files. """
    seqs_dictionary, seqs_dictionary_inverse = make_sequence_dictionary(align_file)
    tree = tabulate_names(tree_file)
    children_stack = get_children_stack(tree, seqs_dictionary, seqs_dictionary_inverse, cutoff )
    remove_sub_nodes(children_stack)
    # split_children_nodes_below_cutoff(children_stack, seqs_dictionary )
    csv = open('results.csv', 'w')
    family_count = 1
    for child in children_stack:
        align_file_name = "%s_align.a2m" % (family_count)
        align_file = open(align_file_name, 'w')
        for terminal in child.get_terminals():
            csv_output = "%s,%s,%s\n" %( family_count, terminal.name.split("_")[1], seqs_dictionary[terminal.name.split("_")[1]][2] )
            csv.write(csv_output)
            a2m_output=">%s:%s\n%s\n" % ( family_count, terminal.name.split("_")[1], seqs_dictionary[terminal.name.split("_")[1]][2] )
            align_file.write(a2m_output)
            if terminal.name.split("_")[1] in seqs_dictionary:
                del  seqs_dictionary[terminal.name.split("_")[1]]
        family_count = family_count + 1
        align_file.close()
    # make_remaining_seqs_family(family_count, seqs_dictionary)
    if len(seqs_dictionary) > 0:
        align_file_name = "%s_align.a2m" % (family_count)
        align_file = open(align_file_name, 'w')
        for remaining_sequence in seqs_dictionary:
            csv_output = "%s,%s,%s\n" %( family_count, remaining_sequence, seqs_dictionary[remaining_sequence][2] )
            csv.write(csv_output)
            a2m_output=">%s:%s\n%s\n" %( family_count, remaining_sequence, seqs_dictionary[remaining_sequence][2] )
            align_file.write(a2m_output)
        align_file.close()
        family_count = family_count + 1
    csv.close()
    print "Program successful.  Results in csv.txt. Number of align files made:%s" % int(family_count-1)


def get_children_stack(tree, seqs_dictionary, seqs_dictionary_inverse, cutoff ):
    """ find highest common clade below seq cutoff and returns stack of children """
    children_stack = []
    for clade in tree.get_terminals():
        terminal_sequence = seqs_dictionary[clade.name.split("_")[1]][2]
        for sequence in seqs_dictionary_inverse:
            sequence_id =  get_seq_id(terminal_sequence[0:20], sequence[0:20])
            #print "terminal_sequence:\n%s\nSequence:\n%s\nid:%s\n" % (terminal_sequence[0:20], sequence[0:20], sequence_id )
            if sequence_id < cutoff:
                second_clade = None
                for find_clade in tree.get_terminals():
                    if seqs_dictionary_inverse[sequence] in find_clade.name: second_clade = find_clade
                common_clade = tree.common_ancestor(clade, second_clade)
                print "seq:%s id:%s highest clade%s" %(sequence_id, cutoff, common_clade.name)
                for find_children in tree.find_clades():
                    if find_children in common_clade:
                        if find_children not in children_stack:
                            children_stack.append(find_children)
    print children_stack
    return children_stack

def remove_sub_nodes(children_stack):
    """ Looks for clades in children stack which are contained within higher clades and removes them"""
    remove_sub_nodes_stack = []
    for child in children_stack:
        for terminal in child.get_terminals():
            for child_duplicate_check in children_stack:
                if terminal in child_duplicate_check.get_terminals() and child != child_duplicate_check:
                    if child not in remove_sub_nodes_stack: remove_sub_nodes_stack.append(child)
    for child in remove_sub_nodes_stack:
        children_stack.remove(child)
    print children_stack
    # todo: need to fix here - may by using collapse or split
    return children_stack

def tabulate_names(tree_file):
    """ Adapted from from BioPython cookbook.  Gives clades a unique name"""
    f = open(tree_file, 'r')
    tree_text = f.read()
    f.close()
    tree = Phylo.read(StringIO(tree_text), "newick")
    names = {}
    for idx, clade in enumerate(tree.find_clades()):
        if clade.name:
            clade.name = '%d_%s' % (idx, clade.name)
        else:
            clade.name = str(idx)
    return tree

def get_seq_id(sequence1, sequence2):
    """Calculate seq id"""
    align_length = 0
    seq_id = 0
    for i, c in enumerate(sequence1 ):
        res1 = sequence1[i]
        res2 = sequence2[i]
        if (res1 == res2):
            seq_id = seq_id + 1
            align_length = align_length + 1
        elif ( (res1 == "-" and res2 != "-") or (res1 != "-" and res2 == "-") or (res1 != "-" and res2 != "-")):
            align_length = align_length + 1

##        if (res1 == res2) and res1 != "-" and res1 != ".":
##            align_length = align_length + 1
##        if (res1 == "-" and res2 != "-") or (res1 != "-" and res2 == "-"):
##            align_length = align_length + 1
    percent_seq_id  = round( (float(seq_id) / align_length)*100, 5)
    return percent_seq_id

def make_sequence_dictionary(align_file):
    """Create sequence dictionarys"""
    f = open(align_file, 'r')
    raw_seqs = f.read()
    f.close()
    seqs_dictionary = {}
    seqs_dictionary_inverse = {}
    raw_seqs_list = raw_seqs.split('>')
    raw_seqs_list.pop(0)
    for seq in raw_seqs_list:
        header_fasta_split = seq.split('\n')
        header_split = header_fasta_split[0].split('|')
        seqs_dictionary[ header_split[1] ] = [ header_split[1], header_split[2], header_fasta_split[1] ]
        seqs_dictionary_inverse [ header_fasta_split[1] ] = header_split[1]
    return seqs_dictionary, seqs_dictionary_inverse

def asserts():
    """basic asserts for program"""
    assert get_seq_id("MSTPP----W","-TTPPPPP-W") == 50


if __name__ == '__main__':
    """command line: kerf.py cutoff align_file tree_file"""
    asserts()
    print get_seq_id("..MLEKHRHLLIKFRQDISRDLLVDDVLPYLQSKLVLDLYDKECIYVEKTSKKKAQRLLD-ILPTKGFDAFEHFFTIL-GEKYPQLVGLLRSGVSDETYGD-AIDS------------NVLLEG--------------GVPQRPDIFTNRARDVHKVRTALRSLKEHDGWVILHGMAGCGKTVLAAEALRSPGLLDQLFPGGVFWIRIGSVSQAKLLMQMQNLCVRLD..--HDHSRAPPRNLEEAKDRLRMLFAHQYPRSLLVLDDLWNAADVKYFDIRCRTLVTTRDASITDSVGGSKVKVRVSEGFSDKEALQILSLW--TRTFDLPPEAIKLVNLCNGSPLAISMIGA-MLKE---HPNRWKYYLNQFQKKKVSKLKTKFAYEYPSLSEAILMGINNLSEEMRTYYIDFALFDDGSKVPAQVLCILWDEELEIVEDLMDELVNKSLAR...KCVQSQDALGISYSIHNLQLAFLQEL-STDKVALHAKIIDRYRDVyggTGGKFDVMEDDHYIHWN--LIRHMLKANMVEDTYTLITSLSWVVSKLKVTGPSDVISDYISVRdSIKAEQELEVLDNFHHFVSTNAHLFVAKMMPDIVQLALQQPSLSQVHTQAMAQASILTKQGKMYIEWCNQRADTLDSRLLITAKVHCGAVYCCKFSSDSSKVVSCGTDNHVKVWDSQSGRQLLSIS-GHGD.VVNCCAFSHDDSRIISCSADQTVKIWDASSGEGVLCFV.GHTQEVFSC.....SFSPDDTKAV--SCSADRT..VKVWDSKTGVCYHVYMAHTDIVRWCCFSPDGGKVASCSDDNTVRIWEASSG---EDLCL---------LDDHT-SSVSFCCFMNQGESVVTVCSNCIKVWSCSSGEQTWVYN.SSSTCL--SLAFSPDDSIVAMGF------SDTTVQLWNSKTFARLGVYKGHKGWAHSVGISKDSSKLVSGSEDETVKIWTID.RDAAKDAAKLRRV------FDAHFGEDSLGVAVA-DSSDRLLIFQGQDGVLTAEGEA--MDSKIHSVTFSPRGEAIAVG-CESGTISIQDRSCKVLQS--FSEH---KGAVRRCYFSQDGQLLLSASDDKTAKVFNLQEARVMYTLAGHTNKLRRCMFFN-NDT.RILTASLDGSLKV.WDAKTGNLEFTC-CHSStdYVLTCDVSHDSQRLISASADCFAKVWDASTGELLLSLGRHPDVVRSISFSPD-NMICTGCDDGIVRIWDSVSGKEVTSC------KKNE---TWIADCHFTRDGRNIVSVSDNIQWWS-KSGQLIQEFFIKGGFAKDVKCSSDFKTFVTVDDTSTLYILNQI----------------------------------------", "..MDAKARNCLLQHKEALEKDIKTSYIMDHMISNGVLTVVEEEKVKSQATQYQRAAALIK-MILNKDNYAYISFYNALLHEGYKDLAGLLHSGLPLVSSSS-GKDT---DGGNTSFVRTVLCEG--------------GVPQRPVIFVTRKKLVSAIQQKLWKLNGEPGWVTIYGMAGCGKSVLAAEAVRDHALLEGCFSGGVHWVSIGKQDKSGLLMKLQNLCTRLG..QEESFSQRLPLNIEEAKDRLRVLMLRKHPRSLLILDDVWDPWVLKAFDNQCQILLTTRDKSVTDSVMGPKYVIPVESGLGKEKGLEILSLFVNMKKEDLPVEAHSIIKECKGSPLVVSLVGA-LLRD---FPNRWAYYLRQLQNKQFKRIRKSSSYDYEALDEAMSISVEMLREDIKDYYTDLSILQKDVKVPTKVLCVLWDLETEEVEDILQEFVNKSLLF...---CNRNGKSFCYYLHDLQVDFLTEKNRSQLQDLHRKMVTQFQRY...HQPHTLSPGQEDCMYWYNFLAYHMASAGMHKELCALMFSLDWIKAKTELVGPAHLIHEFVEYR.HILDEKDCAVCENFQEFLSLNGHLLGRQPFPNIVQLGLCEPETSEVYQQAKLQAKQEVDTGRLYLEWINK--KTIKNLSRLVVRPHTDAVYHACFSQDGQRIASCGADKTLQVFKAETGEKLLDIK-AHED.EVLCCAFSSDDSYIATCSVDKKVKIWDSGTGKLVHTYE.EHSEQVNCC.....HFTNKSNHLLLATGSNDSF..LKLWDLNKKECRNTMFGHTNSVNPCRFSPDDELLASCSADGTLKLWDVRSANEKKSINVKRFFLSSEDPPEDVEVIVKCCSWSADGDRIIVAAKNKVLLLDIHTSGLLTEIH.TGHHSTIQYCDFSPYDHLAVIAL------SQYCVELWNIDSRVKVADCRGHLSWVHGVMFSPDGSSFLTASDDQTIRVWETR.KVCKNSAIVLKQE------IDVVFQENEMMVLAV-DNIRGLQLIAGKTGQIDYLPEA--Q---VSCCCLSPHLEYVAFG-DEEGAIKIIE-LPNNRVFSSGIGH---KKAVRHIQFTADGKTLISSSEDSVIQVWNWQTEEYVF-LQAHQETVKDFRLLRDS--.RLLSWSFDGTVKV.WNVVTGRIERDFTCHQG..TVLSCAISSDATKFSSTSADKTAKIWSFELRSPLHELKGHNSCVRCSAFSLDGILLATGDDNGEIRIWNVSDGQLLHLCAPISIEEGTATHGGWVTDVCFSPDRKMLVSAGGYLKWWNVVTGESSQTFYTNGTNLKKIHVSPDFRTYVTVDNLGILYILQVLE---------------------------------------")
    prog, cutoff, align_file, tree_file = sys.argv
    main(cutoff, align_file, tree_file)
    #main(90, 'bpg087857.a2m', 'bpg087857_nj.nh')
